package com.googlecode.thriftbenchmark.rpc;

import java.util.concurrent.atomic.AtomicInteger;

import org.apache.thrift.TException;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;

import com.googlecode.thriftbenchmark.twitter.Tweet;
import com.googlecode.thriftbenchmark.twitter.TweetFactory;
import com.googlecode.thriftbenchmark.twitter.TwitterService;

/**
 * @author Li Weibo (weibo.leo@gmail.com) //I believe spring-brother
 * @since 2011-12-25 下午04:04:25
 */
public class BenmarkClient {

    private static String transType;
    private static String protocType;
    private static String host;
    private static int port;
    private static int nThreads;
    private static int nRequests;
    private static int nDelayedMs;
    
    
    private static long startTime;
    
    private static AtomicInteger count = new AtomicInteger();
    
    private TTransport getSocketTransport(String host, int port) throws TTransportException {
        TTransport transport = new TSocket(host, port);
        transport.open();
        return transport;
    }
    
    private TTransport getFramedTransport(String host, int port) throws TTransportException {
        TTransport transport = new TFramedTransport(new TSocket(host, port));
        transport.open();
        return transport;
    }
    
    /*private TTransport getNioTransport(String host, int port) throws TTransportException, IOException {
        TTransport transport = new TNonblockingSocket(host, port);
        transport.open();
        return transport;
    }*/
    
    
    
    private class Task implements Runnable {

        @Override
        public void run() {
            
            try {
                TTransport transport = null;
                if ("socket".equals(transType)) {
                    transport = getSocketTransport(host, port);
                } else if ("framed".equals(transType)) {
                    transport = getFramedTransport(host, port);
                } else {
                    System.err.println("Invalid transport type: " + transType);
                    return;
                }
                SimpleTProtocolFactory proFact = new SimpleTProtocolFactory(protocType);
                TProtocol protocol = proFact.getProtocol(transport);
                TwitterService.Client client = new TwitterService.Client(protocol);
                TweetFactory fact = new TweetFactory();
                
                for (int i = 0; i < nRequests; i++) {
                    Tweet t = fact.randomTweet();
                    Tweet rt = client.postAndReturnTweet(t, nDelayedMs);
                    if (rt == null) {
                        System.err.println("ret null");
                        return;
                    }
                    count.incrementAndGet();
                }
                
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        
    }
    
    public static void main(String[] args) throws TException {

        if (args.length != 7) {
            System.err.println("Usage: java " + BenmarkClient.class.getName() + " transport protocol host port nThreads nRequests delayedMs");
            return;
        }
        
        BenmarkClient app = new BenmarkClient();
        
        transType = args[0];
        protocType = args[1];
        host = args[2];
        port = Integer.parseInt(args[3]);
        nThreads = Integer.parseInt(args[4]);
        nRequests = Integer.parseInt(args[5]);
        nDelayedMs = Integer.parseInt(args[6]);
        
        System.out.println("Benchmark start");
        
        Thread[] threads = new Thread[nThreads];
        for (int i = 0; i < threads.length; i++) {
            threads[i] = new Thread(app.new Task());
        }
        
        for (int i = 0; i < threads.length; i++) {
            threads[i].start();
        }
        startTime = System.currentTimeMillis();
        
        startProcessMonitor();
        
        for (int i = 0; i < threads.length; i++) {
            try {
                threads[i].join();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
        
        logQps();
        System.out.println("done");
    }
    
    private static void startProcessMonitor() {
        Thread t = new Thread() {
            public void run() {
                while(true) {
                    try {
                        Thread.sleep(1000 * 10);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                    logQps();
                }
            }
            
        };
        t.setDaemon(true);
        t.start();
    }
    
    private static void logQps() {
        long end = System.currentTimeMillis();
        double cost = (end - startTime) * 1.0 / 1000;
        System.out.println(count.get() + " requests finished in " + cost + "s, qps: " + count.get() / cost);
    }
    
}
